package com.zx.crypto;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.Base64;

public class AESExample {
    private static final String ALGORITHM = "AES";
    private static final String ECB_TRANSFORMATION = "AES/ECB/PKCS5Padding";
    private static final String CBC_TRANSFORMATION = "AES/CBC/PKCS5Padding";

    public static String encryptECB(String plainText, String key) throws Exception {
        SecretKeySpec secretKeySpec = new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), ALGORITHM);
        Cipher cipher = Cipher.getInstance(ECB_TRANSFORMATION);
        cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
        byte[] encryptedBytes = cipher.doFinal(plainText.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(encryptedBytes);
    }

    public static String decryptECB(String cipherText, String key) throws Exception {
        SecretKeySpec secretKeySpec = new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), ALGORITHM);
        Cipher cipher = Cipher.getInstance(ECB_TRANSFORMATION);
        cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
        byte[] decodedBytes = Base64.getDecoder().decode(cipherText);
        byte[] decryptedBytes = cipher.doFinal(decodedBytes);
        return new String(decryptedBytes, StandardCharsets.UTF_8);
    }

    public static String encryptCBC(String plainText, String key) throws Exception {
        SecretKeySpec secretKeySpec = new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), ALGORITHM);
        Cipher cipher = Cipher.getInstance(CBC_TRANSFORMATION);
        byte[] iv = new byte[cipher.getBlockSize()];
        SecureRandom secureRandom = new SecureRandom();
        secureRandom.nextBytes(iv);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec);
        byte[] encryptedBytes = cipher.doFinal(plainText.getBytes(StandardCharsets.UTF_8));
        byte[] encryptedWithIV = new byte[iv.length + encryptedBytes.length];
        System.arraycopy(iv, 0, encryptedWithIV, 0, iv.length);
        System.arraycopy(encryptedBytes, 0, encryptedWithIV, iv.length, encryptedBytes.length);
        return Base64.getEncoder().encodeToString(encryptedWithIV);
    }

    public static String decryptCBC(String cipherText, String key) throws Exception {
        byte[] decodedBytes = Base64.getDecoder().decode(cipherText);
        SecretKeySpec secretKeySpec = new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), ALGORITHM);
        Cipher cipher = Cipher.getInstance(CBC_TRANSFORMATION);
        byte[] iv = new byte[cipher.getBlockSize()];
        System.arraycopy(decodedBytes, 0, iv, 0, iv.length);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
        byte[] encryptedBytes = new byte[decodedBytes.length - iv.length];
        System.arraycopy(decodedBytes, iv.length, encryptedBytes, 0, encryptedBytes.length);
        cipher.init(Cipher.DECRYPT_MODE, secretKeySpec, ivParameterSpec);
        byte[] decryptedBytes = cipher.doFinal(encryptedBytes);
        return new String(decryptedBytes, StandardCharsets.UTF_8);
    }

    public static void main(String[] args) {
        try {
            String plainText = "Hello, AES!";
            // AES 密钥长度必须为 16、24 或 32 字节
            String key = "0123456789abcdef";

            // ECB 模式
            String encryptedECB = encryptECB(plainText, key);
            String decryptedECB = decryptECB(encryptedECB, key);
            System.out.println("ECB Encrypted: " + encryptedECB);
            System.out.println("ECB Decrypted: " + decryptedECB);

            // CBC 模式
            String encryptedCBC = encryptCBC(plainText, key);
            String decryptedCBC = decryptCBC(encryptedCBC, key);
            System.out.println("CBC Encrypted: " + encryptedCBC);
            System.out.println("CBC Decrypted: " + decryptedCBC);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
